package org.rbio.paper.methylglm.annotate;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

import org.broadinstitute.sting.gatk.datasources.reference.ReferenceDataSource;
import org.broadinstitute.sting.utils.GenomeLocParser;
import org.broadinstitute.sting.utils.codecs.refseq.RefSeqCodec;

// prepare annotation files suitable for GATK
public class PrepareAnnotationTables {

	// process
	public static String prepare(String ref_fa, String ucsc_gene,
			String cpg_isl, String meth_file, int tss_up, int tss_down)
			throws Exception {
		// make random prefix
		String out_prefix = System.getProperty("java.io.tmpdir") + "/"
				+ Calendar.getInstance().getTimeInMillis();

		// cpg island table
		SortedTable isl_tab = sort(cpg_isl, 1, 2, ref_fa);
		writeTable(isl_tab, out_prefix + ".isl.gatk.txt", 1, 2, 3, true);

		// cpg table
		SortedTable cpg_tab = sort(meth_file, 0, 1, ref_fa);
		writeTable(cpg_tab, out_prefix + ".cpg.gatk.txt", 0, 1, -1, true);

		// UCSC gene table
		SortedTable ucsc_tab = sort(ucsc_gene, 2, 4, ref_fa);
		writeUCScGene(ucsc_tab, out_prefix + ".ucsc.gatk.txt", ref_fa);

		// make tss table
		TSSTable.buildTSS(out_prefix + ".ucsc.gatk.txt", ref_fa, tss_up,
				tss_down, out_prefix + ".tss.gatk.txt");

		return out_prefix;
	}

	// filter the files based on the reference fasta and add gatk header if
	// needed
	public static SortedTable sort(String ifname, int ref_col, int start_col,
			String fa) throws Exception {
		// get list of references
		ReferenceDataSource data_src = new ReferenceDataSource(new File(fa));

		// filter the file [ref [pos [line]]]
		Map<String, Map<Integer, List<String>>> tab = new TreeMap<String, Map<Integer, List<String>>>();
		BufferedReader ibuf = new BufferedReader(new FileReader(ifname));
		String header = ibuf.readLine();
		String line = null;
		while ((line = ibuf.readLine()) != null) {
			String arr[] = line.split("\t");
			String ref = arr[ref_col];

			// check if ref is in the fasta file
			if (data_src.getReference().getSequenceDictionary()
					.getSequence(ref) == null) {
				continue;
			}

			// add the line to the table
			int pos = Integer.parseInt(arr[start_col]);
			Map<Integer, List<String>> ref_tab = tab.get(ref);
			if (ref_tab == null) {
				ref_tab = new TreeMap<Integer, List<String>>();
				tab.put(ref, ref_tab);
			}
			List<String> pos_lst = ref_tab.get(pos);
			if (pos_lst == null) {
				pos_lst = new ArrayList<String>();
				ref_tab.put(pos, pos_lst);
			}
			pos_lst.add(line);
		}
		ibuf.close();
		return new SortedTable(header, tab);
	}

	// write table
	public static void writeTable(SortedTable sort_tab, String ofname,
			int ref_col, int start_col, int end_col, boolean add_gatk_header)
			throws Exception {
		// write the sorted file
		BufferedWriter obuf = new BufferedWriter(new FileWriter(ofname));
		String header = sort_tab.header;
		header = (add_gatk_header == true) ? "HEADER\t" + header : header;
		obuf.write(header + "\n");
		for (String ref : sort_tab.tab.keySet()) {
			for (Integer pos : sort_tab.tab.get(ref).keySet()) {
				for (String rec : sort_tab.tab.get(ref).get(pos)) {
					if (add_gatk_header == true) {
						String arr[] = rec.split("\t");
						if (end_col >= 0) {
							rec = arr[ref_col] + ":" + arr[start_col] + "-"
									+ arr[end_col] + "\t" + rec;
						} else {
							rec = arr[ref_col] + ":" + arr[start_col] + "\t"
									+ rec;
						}
					}
					obuf.write(rec + "\n");
				}
			}
		}
		obuf.close();
	}

	// write the ucsc file
	public static void writeUCScGene(SortedTable sort_tab, String ofname,
			String ref_fa) throws Exception {
		// prepare the UCSC gene codec
		ReferenceDataSource data_src = new ReferenceDataSource(new File(ref_fa));
		GenomeLocParser parser = new GenomeLocParser(data_src.getReference());
		RefSeqCodec codec = new RefSeqCodec();
		codec.setGenomeLocParser(parser);

		// parse each line
		BufferedWriter obuf = new BufferedWriter(new FileWriter(ofname));
		String header = sort_tab.header;
		obuf.write(header + "\n");
		for (String ref : sort_tab.tab.keySet()) {
			for (Integer pos : sort_tab.tab.get(ref).keySet()) {
				for (String rec : sort_tab.tab.get(ref).get(pos)) {
					// check for non coding tx
					if (codec.decode(rec) != null) {
						obuf.write(rec + "\n");
					}
				}
			}
		}
		obuf.close();
	}

	// helper class to for sorted table
	public static class SortedTable {
		String header;
		Map<String, Map<Integer, List<String>>> tab;

		// constructor
		public SortedTable(String header,
				Map<String, Map<Integer, List<String>>> tab) {
			this.header = header;
			this.tab = tab;
		}
	}
}
